package com.imchenio.springscheduleexample.instance03;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.imchenio.springscheduleexample.common.mapper.TriggerTaskDAO;
import com.imchenio.springscheduleexample.common.mapper.TriggerTaskDO;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.config.BeanPostProcessor;
import org.springframework.scheduling.annotation.SchedulingConfigurer;
import org.springframework.scheduling.config.ScheduledTaskRegistrar;
import org.springframework.scheduling.config.TriggerTask;
import org.springframework.scheduling.support.CronTrigger;
import org.springframework.scheduling.support.ScheduledMethodRunnable;

import java.lang.annotation.Annotation;
import java.lang.reflect.Method;
import java.util.Collections;
import java.util.HashSet;
import java.util.Set;
import java.util.stream.Stream;

@Slf4j
public class EnhanceSchedulingConfig implements BeanPostProcessor, SchedulingConfigurer {

    private static final Set<TriggerTask> scheduleTasks = Collections.synchronizedSet(new HashSet<>());

    private volatile ScheduledTaskRegistrar registrar;

    @Override
    public void configureTasks(ScheduledTaskRegistrar registrar) {
        this.registrar = registrar;
        scheduleTasks.forEach(triggerTask -> this.registrar.addTriggerTask(triggerTask));
        log.info("register {} schedules", scheduleTasks.size());
        scheduleTasks.clear();
    }

    private final TriggerTaskDAO triggerTaskDAO;

    public EnhanceSchedulingConfig(TriggerTaskDAO triggerTaskDAO) {
        this.triggerTaskDAO = triggerTaskDAO;
    }

    @Override
    public Object postProcessAfterInitialization(Object bean, String beanName) throws BeansException {
        registerTaskFromBean(bean);
        return null;
    }

    public void registerTaskFromBean(Object bean) {
        Class<?> clazz = bean.getClass();
        Stream.of(clazz.getMethods()).forEach(method -> {
            registerTaskFromMethod(bean, method);
        });
    }

    public void registerTaskFromMethod(Object object, Method method) {
        for (Annotation annotation : method.getAnnotations()) {
            if (annotation instanceof EnhanceScheduled) {
                ScheduledMethodRunnable runnable = new ScheduledMethodRunnable(object, method);
                register(runnable, (EnhanceScheduled) annotation);
            }
        }
    }

    public void register(ScheduledMethodRunnable runnable, EnhanceScheduled annotation) {
        if (annotation.id() <= 0) {
            return;
        }
        TriggerTask triggerTask = new TriggerTask(runnable, new EnhanceCronTrigger(annotation, triggerTaskDAO));

        if (registrar == null) {
            log.info("add schedule, id = {}", annotation.id());
            scheduleTasks.add(triggerTask);
        } else {
            log.info("register schedule, id = {}", annotation.id());
            registrar.scheduleTriggerTask(triggerTask);
        }
    }
}
